/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.core;



import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.transforms.Materializations;

import com.bff.gaia.unified.sdk.transforms.Materializations.MultimapView;

import com.bff.gaia.unified.sdk.values.KV;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ArrayListMultimap;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.Multimap;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.Multimaps;



/** An in-memory representation of {@link MultimapView}. */

public class InMemoryMultimapSideInputView<K, V> implements Materializations.MultimapView<K, V> {



  /**

   * Creates a {@link MultimapView} from the provided values. The provided {@link Coder} is used to

   * guarantee structural equality for keys instead of assuming Java object equality.

   */

  public static <K, V> MultimapView<K, V> fromIterable(

      Coder<K> keyCoder, Iterable<KV<K, V>> values) {

    // We specifically use an array list multimap to allow for:

    //  * null keys

    //  * null values

    //  * duplicate values

    Multimap<Object, Object> multimap = ArrayListMultimap.create();

    for (KV<K, V> value : values) {

      multimap.put(keyCoder.structuralValue(value.getKey()), value.getValue());

    }

    return new InMemoryMultimapSideInputView(keyCoder, Multimaps.unmodifiableMultimap(multimap));

  }



  private final Coder<K> keyCoder;

  private final Multimap<Object, V> structuralKeyToValuesMap;



  private InMemoryMultimapSideInputView(Coder<K> keyCoder, Multimap<Object, V> data) {

    this.keyCoder = keyCoder;

    this.structuralKeyToValuesMap = data;

  }



  @Override

  public Iterable<V> get(K k) {

    return structuralKeyToValuesMap.get(keyCoder.structuralValue(k));

  }

}